from typing import Dict, List, Optional, Any, Union
from pydantic import BaseModel, Field
from enum import Enum
from app.schemas.facenet_base_schema import TaskStatus

class NodeType(str, Enum):
    """节点类型枚举"""
    DEFAULT = "default"
    PERSON = "person"
    COMPANY = "company"
    PRODUCT = "product"
    # 可以根据实际需求添加更多类型

class EdgeType(str, Enum):
    """边类型枚举"""
    DEFAULT = "default"
    KNOWS = "knows"
    BELONGS_TO = "belongs_to"
    PRODUCES = "produces"
    # 可以根据实际需求添加更多类型

class GraphNode(BaseModel):
    """图节点"""
    id: str
    type: NodeType
    properties: Dict[str, Any] = {}

class GraphEdge(BaseModel):
    """图边"""
    source: str
    target: str
    type: EdgeType = EdgeType.DEFAULT
    properties: Dict[str, Any] = {}

class Graph(BaseModel):
    """图结构"""
    nodes: List[GraphNode]
    edges: List[GraphEdge]

class TaskType(str, Enum):
    NODE_CLASSIFICATION = "node_classification"
    COMMUNITY_DETECTION = "community_detection"
    # 可以添加其他任务类型

class HGCNParams(BaseModel):
    """HGCN算法参数"""
    graph: Graph
    hidden_dims: List[int] = [64, 32]
    learning_rate: float = 0.01
    epochs: int = 100
    task: TaskType = TaskType.NODE_CLASSIFICATION
    labels: Optional[Dict[str, int]] = None

class HGCNRequest(BaseModel):
    """HGCN请求"""
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: HGCNParams

class NodeEmbedding(BaseModel):
    """节点嵌入"""
    node_id: str
    embedding: List[float]

class PredictionResult(BaseModel):
    """预测结果"""
    node_id: str
    predicted_label: int

class HGCNResult(BaseModel):
    """HGCN算法结果"""
    embeddings: List[NodeEmbedding]
    predictions: Optional[List[PredictionResult]] = None

class HGCNResponse(BaseModel):
    """HGCN响应"""
    task_id: str
    task_status: TaskStatus
    task_progress: float = 100.0
    output_params: HGCNResult
    metrics: List[Dict[str, Any]] = []
    error_message: Optional[str] = None

class HGCNMiddleResponse(BaseModel):
    """HGCN中间响应"""
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: float
    input_params: HGCNParams
    metrics: List[Dict[str, Any]] = []
    error_message: Optional[str] = None 